#include <ros/ros.h>
#include <sensor_msgs/PointCloud2.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <nav_msgs/Odometry.h>

#include <pcl_ros/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl/registration/icp.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/registration/correspondence_types.h>

#include <Eigen/Dense>
#include <thread>
#include <mutex>

// typedefs
using PointT = pcl::PointXYZI;
using PointCloudT = pcl::PointCloud<PointT>;

ros::Publisher pub_pc_in_map, pub_submap, pub_map_to_odom;
PointCloudT::Ptr global_map(new PointCloudT);
PointCloudT::Ptr cur_scan(new PointCloudT);
nav_msgs::Odometry cur_odom;
Eigen::Matrix4f T_map_to_odom = Eigen::Matrix4f::Identity();
Eigen::Matrix4f initial_pose = Eigen::Matrix4f::Identity();

bool global_map_initialized = false;
bool initial_pose_received = false;
bool initialized = false;

std::mutex data_mutex;

const float MAP_VOXEL_SIZE = 0.3;
const float SCAN_VOXEL_SIZE = 0.1;
const double FREQ_LOCALIZATION = 0.5;
const float LOCALIZATION_TH = 0.90;

const float FOV = 6.2831853;
const float FOV_FAR = 30.0;

// --- Utility functions ---
Eigen::Matrix4f poseToMat(const geometry_msgs::PoseWithCovarianceStamped &msg)
{
    Eigen::Matrix4f mat = Eigen::Matrix4f::Identity();
    const auto &p = msg.pose.pose.position;
    const auto &q = msg.pose.pose.orientation;
    Eigen::Quaternionf quat(q.w, q.x, q.y, q.z);
    Eigen::Matrix3f R = quat.toRotationMatrix();
    mat.block<3, 3>(0, 0) = R;
    mat(0, 3) = p.x;
    mat(1, 3) = p.y;
    mat(2, 3) = p.z;
    return mat;
}

Eigen::Matrix4f inverseSE3(const Eigen::Matrix4f &T)
{
    Eigen::Matrix4f inv = Eigen::Matrix4f::Identity();
    inv.block<3, 3>(0, 0) = T.block<3, 3>(0, 0).transpose();
    inv.block<3, 1>(0, 3) = -inv.block<3, 3>(0, 0) * T.block<3, 1>(0, 3);
    return inv;
}

PointCloudT::Ptr voxelDownSample(PointCloudT::Ptr cloud, float voxel_size)
{
    pcl::VoxelGrid<PointT> vg;
    PointCloudT::Ptr cloud_filtered(new PointCloudT);
    vg.setInputCloud(cloud);
    vg.setLeafSize(voxel_size, voxel_size, voxel_size);
    vg.filter(*cloud_filtered);
    return cloud_filtered;
}

Eigen::Matrix4f runICP(PointCloudT::Ptr src, PointCloudT::Ptr tgt, Eigen::Matrix4f init_guess, float voxel_scale, float &fitness_score)
{

    if (!src || src->empty() || !tgt || tgt->empty())
    {
        ROS_WARN(" >>> [ ICP ] Input point cloud is empty.");
        fitness_score = 0.0f;
        return Eigen::Matrix4f::Identity();
    }

    ROS_INFO(" >>> [ localization ] Running ICP...");
    // 降采样处理
    auto src_filtered = voxelDownSample(src, SCAN_VOXEL_SIZE * voxel_scale);
    auto tgt_filtered = voxelDownSample(tgt, MAP_VOXEL_SIZE * voxel_scale);

    // 配置 ICP 参数
    pcl::IterativeClosestPoint<PointT, PointT> icp;
    icp.setInputSource(src_filtered);
    icp.setInputTarget(tgt_filtered);
    icp.setMaximumIterations(50);
    icp.setMaxCorrespondenceDistance(1.5f * voxel_scale);
    icp.setRANSACOutlierRejectionThreshold(1.5f * voxel_scale);

    // 执行 ICP 配准
    PointCloudT aligned;
    icp.align(aligned, init_guess);

    // 手动计算内点比例（fitness_score）
    int inlier_count = 0;
    float threshold_sq = (1.5f * voxel_scale) * (1.5f * voxel_scale); // 距离阈值平方
    pcl::search::KdTree<PointT> kdtree;
    kdtree.setInputCloud(tgt_filtered);

    std::vector<int> indices(1);
    std::vector<float> distances_sq(1);
    for (const auto &point : *src_filtered)
    {
        kdtree.nearestKSearch(point, 1, indices, distances_sq);
        if (distances_sq[0] <= threshold_sq)
        {
            inlier_count++;
        }
    }

    fitness_score = static_cast<float>(inlier_count) / static_cast<float>(src_filtered->size());

    return icp.hasConverged() ? icp.getFinalTransformation() : Eigen::Matrix4f::Identity();
}

PointCloudT::Ptr cropGlobalMapInFOV(const Eigen::Matrix4f &pose_est, const nav_msgs::Odometry &odom)
{
    ROS_INFO(" >>> [ localization ] Running crop global map in FOV...");

    // 检查 global_map 是否已初始化
    if (!global_map || global_map->empty())
    {
        ROS_WARN(" >>> [ localization ] Global map is null or empty. Skipping FOV cropping.");
        return PointCloudT::Ptr();
    }

    // 检查 odom 是否包含有效数据
    const auto &p = odom.pose.pose.position;
    const auto &q = odom.pose.pose.orientation;
    if (std::isnan(p.x) || std::isnan(p.y) || std::isnan(p.z) ||
        std::isnan(q.w) || std::isnan(q.x) || std::isnan(q.y) || std::isnan(q.z))
    {
        ROS_WARN(" >>> [ localization ] Odometry data contains NaN values. Skipping FOV cropping.");
        return PointCloudT::Ptr();
    }

    // 构造 T_odom_to_base（里程计到当前坐标系的变换）
    Eigen::Matrix4f T_odom_to_base = Eigen::Matrix4f::Identity();
    Eigen::Quaternionf quat(q.w, q.x, q.y, q.z);
    T_odom_to_base.block<3, 3>(0, 0) = quat.toRotationMatrix();
    T_odom_to_base.block<3, 1>(0, 3) << p.x, p.y, p.z;

    // 计算 T_map_to_base（全局地图到当前坐标系的变换）
    Eigen::Matrix4f T_map_to_base = pose_est * T_odom_to_base;
    Eigen::Matrix4f T_base_to_map = inverseSE3(T_map_to_base);

    // 检查变换矩阵是否合法
    if (T_base_to_map.hasNaN())
    {
        ROS_WARN(" >>> [ localization ] Transformation matrix contains NaN. Skipping FOV cropping.");
        return PointCloudT::Ptr();
    }

    // 裁剪子地图
    PointCloudT::Ptr submap(new PointCloudT);
    for (const auto &pt : global_map->points)
    {
        Eigen::Vector4f p_global(pt.x, pt.y, pt.z, 1.0);
        Eigen::Vector4f p_base = T_base_to_map * p_global;

        // 判断点是否在视野范围内（距离和角度）
        float dist = std::sqrt(p_base(0) * p_base(0) + p_base(1) * p_base(1));
        float angle = std::atan2(p_base(1), p_base(0));
        if (std::abs(angle) < FOV / 2.0 && dist < FOV_FAR)
        {
            submap->points.push_back(pt);
        }
    }

    // 检查裁剪后的子图是否为空
    if (submap->empty())
    {
        ROS_WARN(" >>> [ localization ] Submap is empty after FOV cropping.");
        return PointCloudT::Ptr();
    }

    // 设置子地图头信息并发布
    submap->header.frame_id = "map";
    pcl_conversions::toPCL(ros::Time::now(), submap->header.stamp);
    pub_submap.publish(*submap);

    ROS_INFO(" >>> [ localization ] Cropping done. Submap size: %zu points", submap->size());
    return submap;
}

bool globalLocalization(const Eigen::Matrix4f &pose_guess)
{
    std::lock_guard<std::mutex> lock(data_mutex);
    if (cur_scan->empty())
        return false;

    auto submap = cropGlobalMapInFOV(pose_guess, cur_odom);

    if (!submap || submap->empty())
    {
        ROS_WARN(" >>> [ localization ] Submap is empty.");
        return false;
    }
    // 多尺度 ICP 配准（粗略 → 精细）
    float fitness;
    auto T_init = runICP(cur_scan, submap, pose_guess, 5.0, fitness);
    auto T_final = runICP(cur_scan, submap, T_init, 1.0, fitness);

    // 判断配准是否成功
    if (fitness > LOCALIZATION_TH)
    {
        T_map_to_odom = T_final;

        // 构造并发布 Odometry 消息
        nav_msgs::Odometry odom_msg;
        Eigen::Vector3f trans = T_map_to_odom.block<3, 1>(0, 3);
        Eigen::Quaternionf rot(T_map_to_odom.block<3, 3>(0, 0));
        odom_msg.pose.pose.position.x = trans.x();
        odom_msg.pose.pose.position.y = trans.y();
        odom_msg.pose.pose.position.z = trans.z();
        odom_msg.pose.pose.orientation.x = rot.x();
        odom_msg.pose.pose.orientation.y = rot.y();
        odom_msg.pose.pose.orientation.z = rot.z();
        odom_msg.pose.pose.orientation.w = rot.w();
        odom_msg.header.frame_id = "map";
        odom_msg.header.stamp = ros::Time::now();
        pub_map_to_odom.publish(odom_msg);
        ROS_INFO(" >>> [ localization ] Successfully ICP match. Fitness = %.3f", fitness);
        return true;
    }

    ROS_WARN(" >>> [ localization ] ICP match failed. Fitness = %.3f", fitness);
    return false;
}

// --- Callbacks ---
void cbScan(const sensor_msgs::PointCloud2ConstPtr &msg)
{
    std::lock_guard<std::mutex> lock(data_mutex);
    pcl::fromROSMsg(*msg, *cur_scan);
    pub_pc_in_map.publish(*msg);
}

void cbOdom(const nav_msgs::OdometryConstPtr &msg)
{
    std::lock_guard<std::mutex> lock(data_mutex);
    cur_odom = *msg;
}

void cbMap(const sensor_msgs::PointCloud2ConstPtr &msg)
{
    std::lock_guard<std::mutex> lock(data_mutex); // 加锁
    pcl::fromROSMsg(*msg, *global_map);
    global_map = voxelDownSample(global_map, MAP_VOXEL_SIZE);
    global_map_initialized = true;
    ROS_INFO(" >>> [ localization ] -- Global map received.");
}

void cbInitialPose(const geometry_msgs::PoseWithCovarianceStampedConstPtr &msg)
{
    initial_pose = poseToMat(*msg);
    T_map_to_odom = initial_pose;
    initial_pose_received = true;
    ROS_INFO(" >>> [ localization ] -- Initial pose received.");
}

void localizationThread()
{
    ros::Rate rate(FREQ_LOCALIZATION);
    while (ros::ok())
    {
        if (global_map_initialized && initial_pose_received)
            globalLocalization(T_map_to_odom);
        rate.sleep();
    }
}

// --- Main ---
int main(int argc, char **argv)
{
    ros::init(argc, argv, "fast_lio_localization_cpp");
    ros::NodeHandle nh;

    pub_pc_in_map = nh.advertise<sensor_msgs::PointCloud2>("/cur_scan_in_map", 1);
    pub_submap = nh.advertise<sensor_msgs::PointCloud2>("/submap", 1);
    pub_map_to_odom = nh.advertise<nav_msgs::Odometry>("/map_to_odom", 1);

    ros::Subscriber sub_scan = nh.subscribe("/cloud_registered", 1, cbScan);
    ros::Subscriber sub_odom = nh.subscribe("/Odometry", 1, cbOdom);
    ros::Subscriber sub_map = nh.subscribe("/map3d", 1, cbMap);
    ros::Subscriber sub_init = nh.subscribe("/initialpose", 1, cbInitialPose);

    ROS_WARN(" >>> [ localization ] Waiting for map and initial pose...");

    while (ros::ok() && (!global_map_initialized || !initial_pose_received))
    {
        ros::Duration(0.2).sleep();
        ros::spinOnce();
    }

    if (!cur_scan->empty())
    {
        initialized = globalLocalization(initial_pose);
    }
    else
    {
        ROS_WARN(" >>> [ localization ] First scan not received.");
    }

    std::thread t(localizationThread);
    ros::spin();
    t.join();
    return 0;
}
